# coding = utf-8
import math
import time
from datetime import datetime

import tensorflow as tf

import readdata as rd

batch_size = 10
num_batches = 100


def print_activations(t):
    '''打印每一个卷积层或池化层输出tensor的尺寸
    t:tensor t.op.name：tensor的名称 ;
    t.get-shape.as_list()：tensor尺寸'''
    print(t.op.name, '', t.get_shape().as_list())


def interence(images):
    '''input: images; return: 最后一层pool5及parameters
    '''
    parameters = []

    with tf.name_scope('conv1') as scope:
        # 定义第一个卷积层，卷积核尺寸为11x11，颜色通道为3，卷积核数量为64
        kernel = tf.Variable(tf.truncated_normal([11, 11, 3, 64],
                                                 dtype=tf.float32, stddev=1e-1, name='weights'))
        # 对输入的images进行卷积操作，strides步长设置为4x4
        conv = tf.nn.conv2d(images, kernel, [1, 4, 4, 1], padding='SAME')
        # biases全部初始化为0
        biases = tf.Variable(tf.constant(0.0, shape=[64], dtype=tf.float32),
                             trainable=True, name='biases')
        # 将卷积结果conv和biases
        bias = tf.nn.bias_add(conv, biases)
        # rele对结果进行非线性处理
        conv1 = tf.nn.relu(bias, name=scope)
        print_activations(conv1)
        # 将这一层的参数kernel和biases添加到parameters
        parameters += [kernel, biases]
    # LRN层，depth_radius=4，等都是AlexNet论文中推荐值
    lrn1 = tf.nn.lrn(conv1, 4, bias=1.0, alpha=0.001 / 9, beta=0.75, name='lrn1')
    pool1 = tf.nn.max_pool(lrn1, ksize=[1, 3, 3, 1], strides=[1, 2, 2, 1],
                           padding='VALID', name='pool1')

    print_activations(pool1)

    with tf.name_scope('conv2') as scope:
        # 定义第二个卷积层，卷积核尺寸为5x5，输入通道为64，卷积核数量为192
        kernel = tf.Variable(tf.truncated_normal([5, 5, 64, 192],
                                                 dtype=tf.float32, stddev=1e-1, name='weights'))
        # 对输入的images进行卷积操作，strides步长设置为1x1
        conv = tf.nn.conv2d(pool1, kernel, [1, 1, 1, 1], padding='SAME')
        # biases全部初始化为0
        biases = tf.Variable(tf.constant(0.0, shape=[192], dtype=tf.float32),
                             trainable=True, name='biases')
        # 将卷积结果conv和biases
        bias = tf.nn.bias_add(conv, biases)
        # rele对结果进行非线性处理
        conv2 = tf.nn.relu(bias, name=scope)

        # 将这一层的参数kernel和biases添加到parameters
        parameters += [kernel, biases]

    print_activations(conv2)
    # LRN层，depth_radius=4，等都是AlexNet论文中推荐值
    lrn2 = tf.nn.lrn(conv2, 4, bias=1.0, alpha=0.001 / 9, beta=0.75, name='lrn2')
    pool2 = tf.nn.max_pool(lrn2, ksize=[1, 3, 3, 1], strides=[1, 2, 2, 1],
                           padding='VALID', name='pool2')

    print_activations(pool2)

    with tf.name_scope('conv3') as scope:
        # 定义第三个卷积层，卷积核尺寸为5x5，输入通道为64，卷积核数量为192
        kernel = tf.Variable(tf.truncated_normal([3, 3, 192, 384],
                                                 dtype=tf.float32, stddev=1e-1, name='weights'))
        # 对输入的images进行卷积操作，strides步长设置为1x1
        conv = tf.nn.conv2d(pool2, kernel, [1, 1, 1, 1], padding='SAME')
        # biases全部初始化为0
        biases = tf.Variable(tf.constant(0.0, shape=[384], dtype=tf.float32),
                             trainable=True, name='biases')
        # 将卷积结果conv和biases
        bias = tf.nn.bias_add(conv, biases)
        # rele对结果进行非线性处理
        conv3 = tf.nn.relu(bias, name=scope)

        # 将这一层的参数kernel和biases添加到parameters
        parameters += [kernel, biases]

        print_activations(conv3)

    with tf.name_scope('conv4') as scope:
        # 定义第四个卷积层，卷积核尺寸为3x3，输入通道为384，卷积核数量为256
        kernel = tf.Variable(tf.truncated_normal([3, 3, 384, 256],
                                                 dtype=tf.float32, stddev=1e-1, name='weights'))
        # 对输入的images进行卷积操作，strides步长设置为1x1
        conv = tf.nn.conv2d(conv3, kernel, [1, 1, 1, 1], padding='SAME')
        # biases全部初始化为0
        biases = tf.Variable(tf.constant(0.0, shape=[256], dtype=tf.float32),
                             trainable=True, name='biases')
        # 将卷积结果conv和biases
        bias = tf.nn.bias_add(conv, biases)
        # rele对结果进行非线性处理
        conv4 = tf.nn.relu(bias, name=scope)

        # 将这一层的参数kernel和biases添加到parameters
        parameters += [kernel, biases]

        print_activations(conv4)

    with tf.name_scope('conv5') as scope:
        # 定义第五个卷积层，卷积核尺寸为3x3，输入通道为256，卷积核数量为256
        kernel = tf.Variable(tf.truncated_normal([3, 3, 256, 256],
                                                 dtype=tf.float32, stddev=1e-1, name='weights'))
        # 对输入的images进行卷积操作，strides步长设置为1x1
        conv = tf.nn.conv2d(conv4, kernel, [1, 1, 1, 1], padding='SAME')
        # biases全部初始化为0
        biases = tf.Variable(tf.constant(0.0, shape=[256], dtype=tf.float32),
                             trainable=True, name='biases')
        # 将卷积结果conv和biases
        bias = tf.nn.bias_add(conv, biases)
        # rele对结果进行非线性处理
        conv5 = tf.nn.relu(bias, name=scope)

        # 将这一层的参数kernel和biases添加到parameters
        parameters += [kernel, biases]

        print_activations(conv5)

    pool5 = tf.nn.max_pool(conv5, ksize=[1, 3, 3, 1], strides=[1, 2, 2, 1],
                           padding='VALID', name='pool5')
    print_activations(pool5)
    return pool5, parameters


def time_tensorflow_run(session, target, info_string):
    '''评估AlexNet每轮计算时间
    target:评测的运算算子
    info_string：评测的名称'''
    num_steps_burn_in = 10  # 预热轮数，给程序热身
    total_duration = 0.0  # 总时间
    total_duration_squared = 0.0  # 计算方差

    for i in range(num_batches + num_steps_burn_in):
        print(i)

        start_time = time.time()
        print(start_time)
        _ = session.run(target)
        duration = time.time() - start_time
        # 在初始热身的num_steps_burn_in次迭代后每10轮显示当前迭代所需要的时间
        if i >= num_steps_burn_in:
            if not i % 2:
                print('%s: step %d, duration = %.3f' %
                      (datetime.now(), i - num_steps_burn_in, duration))
            total_duration += duration
            total_duration_squared += duration * duration
        mn = total_duration / num_batches  # 每轮迭代平均耗时
        vr = total_duration_squared / num_batches - mn * mn
        # 平均耗时标准差
        sd = math.sqrt(vr)
        print(time.time())
        print('%s: %s across %d steps, %.3f +/-%.3f sec/batch' %
              (datetime.now(), info_string, num_batches, mn, sd))


def run_benchmark():
    with tf.Graph().as_default():
        image_size = 227
        '''batch_size:每轮迭代样本数
        image_size:图片尺寸
        3：图片颜色通道数'''
        tfrecords_file = 'F:\\001-python\\catvsdogtrain227.tfrecords'
        img, label = rd.read_and_decode(tfrecords_file, batch_size=batch_size)
        img = tf.cast(img, dtype=tf.float32)
        label = tf.cast(label, dtype=tf.int64)
        # img_batch, label_batch = tf.train.shuffle_batch([img, label],
        #                                                 batch_size=10,
        #                                                 capacity=200,  # capacity是队列的长度
        #                                                 min_after_dequeue=100)  # min_after_dequeue是出队后，队列至少剩下min_after_dequeue个数据
        # images = img

        pool5, parameters = interence(img)

        init = tf.global_variables_initializer()
        sess = tf.Session()
        sess.run(init)

        print("开始前向计算")
        threads = tf.train.start_queue_runners(sess=sess)
        for i in range(50):
            val, l = sess.run([img, label])
            # l = to_categorical(l, 12)
            print(val.shape, l)
        # time_tensorflow_run(sess, pool5, "Forward")

        # objective = tf.nn.l2_loss(pool5)
        # grad = tf.gradients(objective, parameters)
        # time_tensorflow_run(sess, grad, "Forward-backward")


run_benchmark()
